{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0AzKKWkk5pon",
        "outputId": "271ccfd6-aa6e-4b1f-d68c-e55b85e59f53"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "start on cuda:0 device.\n",
            "current_dir: /content\n"
          ]
        }
      ],
      "source": [
        "import time\n",
        "import torch\n",
        "import torchvision.datasets\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "import os\n",
        "\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print(\"start on {} device.\".format(device))\n",
        "\n",
        "print(\"current_dir:\", os.getcwd())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZdlXI53KxLHU"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xZltLUllu-LU"
      },
      "outputs": [],
      "source": [
        "#from google.colab import drive\n",
        "#drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m8CiSnoGwrhO",
        "outputId": "8060d17b-562b-4da8-b7f3-78eebd9f78fe"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Files already downloaded and verified\n",
            "训练数据集长度: 50000\n",
            "测试数据集长度: 50000\n",
            "第0轮训练用时：10.171447277069092 s\n",
            "第1轮训练用时：8.9092116355896 s\n",
            "第2轮训练用时：7.872986078262329 s\n",
            "第3轮训练用时：7.256194353103638 s\n",
            "第4轮训练用时：7.852960586547852 s\n",
            "第5轮训练用时：8.463770151138306 s\n",
            "第6轮训练用时：8.685554265975952 s\n",
            "第7轮训练用时：7.379823923110962 s\n",
            "第8轮训练用时：7.482641220092773 s\n",
            "第9轮训练用时：8.358267784118652 s\n",
            "第10轮训练用时：8.274250507354736 s\n",
            "第11轮训练用时：7.803637504577637 s\n",
            "第12轮训练用时：7.2418177127838135 s\n",
            "第13轮训练用时：7.839684963226318 s\n",
            "第14轮训练用时：8.205397605895996 s\n",
            "第15轮训练用时：8.188384056091309 s\n",
            "第16轮训练用时：7.1404547691345215 s\n",
            "第17轮训练用时：7.488996505737305 s\n",
            "第18轮训练用时：8.230464935302734 s\n",
            "第19轮训练用时：8.20580792427063 s\n",
            "第20轮训练用时：7.913569211959839 s\n",
            "第21轮训练用时：7.111084699630737 s\n",
            "第22轮训练用时：7.638668060302734 s\n",
            "第23轮训练用时：8.170094013214111 s\n",
            "第24轮训练用时：8.308802604675293 s\n",
            "第25轮训练用时：7.432401418685913 s\n",
            "第26轮训练用时：7.255666494369507 s\n",
            "第27轮训练用时：8.180928468704224 s\n",
            "第28轮训练用时：8.278382778167725 s\n",
            "第29轮训练用时：8.319466590881348 s\n",
            "第30轮训练用时：7.2851722240448 s\n",
            "第31轮训练用时：8.119832277297974 s\n",
            "第32轮训练用时：8.595227003097534 s\n",
            "第33轮训练用时：8.357483386993408 s\n",
            "第34轮训练用时：7.3174614906311035 s\n",
            "第35轮训练用时：7.544212341308594 s\n",
            "第36轮训练用时：8.499500036239624 s\n",
            "第37轮训练用时：8.322861909866333 s\n",
            "第38轮训练用时：7.758252382278442 s\n",
            "第39轮训练用时：7.27590537071228 s\n",
            "第40轮训练用时：8.017323017120361 s\n",
            "第41轮训练用时：8.40778636932373 s\n",
            "第42轮训练用时：8.386476516723633 s\n",
            "第43轮训练用时：7.195592641830444 s\n",
            "第44轮训练用时：7.105921745300293 s\n",
            "第45轮训练用时：8.497684478759766 s\n",
            "第46轮训练用时：8.29164171218872 s\n",
            "第47轮训练用时：7.639237403869629 s\n",
            "第48轮训练用时：7.067152976989746 s\n",
            "第49轮训练用时：7.728373050689697 s\n"
          ]
        }
      ],
      "source": [
        "\n",
        "# 训练的轮数\n",
        "epoch = 50\n",
        "batch_size = 128\n",
        "dataset_path = \"./cifar10_dataset\"\n",
        "writer = SummaryWriter(\"./logs\")\n",
        "train_data = torchvision.datasets.CIFAR10(dataset_path, train=True, download=True,\n",
        "                                          transform=torchvision.transforms.ToTensor())\n",
        "train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=False)\n",
        "\n",
        "test_data = torchvision.datasets.CIFAR10(dataset_path, train=False, download=True,\n",
        "                                         transform=torchvision.transforms.ToTensor())\n",
        "test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=False)\n",
        "\n",
        "\n",
        "class MyMod(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(MyMod, self).__init__()\n",
        "        self.model = nn.Sequential(\n",
        "            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),\n",
        "            nn.MaxPool2d(2),\n",
        "            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),\n",
        "            nn.MaxPool2d(2),\n",
        "            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),\n",
        "            nn.MaxPool2d(2),\n",
        "            nn.Flatten(),\n",
        "            nn.Linear(in_features=1024, out_features=64),\n",
        "            nn.Linear(in_features=64, out_features=10)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        output = self.model(x)\n",
        "        return output\n",
        "\n",
        "\n",
        "train_data_size = len(train_data)\n",
        "test_data_size = len(test_data)\n",
        "print(\"训练数据集长度: {}\".format(train_data_size))\n",
        "print(\"测试数据集长度: {}\".format(test_data_size))\n",
        "# 创建网络模型\n",
        "mymod = MyMod()\n",
        "mymod = mymod.to(device)\n",
        "# 损失函数\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "loss_fn = loss_fn.to(device)\n",
        "\n",
        "# 优化器\n",
        "# learning_rate = 0.01\n",
        "# 1e-2=1 x (10)^(-2) = 1 /100 = 0.01\n",
        "learning_rate = 1e-2\n",
        "optimizer = torch.optim.SGD(mymod.parameters(), lr=learning_rate)\n",
        "\n",
        "total_train_step = 0\n",
        "total_test_step = 0\n",
        "test_acc_step = 0\n",
        "for i in range(epoch):\n",
        "    mymod.train()\n",
        "    start_time = time.time()\n",
        "    for data in train_dataloader:\n",
        "        imgs, targets = data\n",
        "        imgs = imgs.to(device)\n",
        "        targets = targets.to(device)\n",
        "        outputs = mymod(imgs)\n",
        "\n",
        "        loss = loss_fn(outputs, targets)\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if total_train_step % 1000 == 0:\n",
        "            writer.add_scalar(\"train_loss\", loss.item(), total_train_step)\n",
        "        total_train_step += 1\n",
        "\n",
        "    end_time = time.time()\n",
        "    print(\"第{}轮训练用时：{} s\".format(i, end_time - start_time))\n",
        "    mymod.eval()\n",
        "    total_test_loss = 0\n",
        "    total_accuracy = 0\n",
        "    # 去掉梯度\n",
        "    with torch.no_grad():\n",
        "        for data in test_dataloader:\n",
        "            imgs, targets = data\n",
        "            imgs = imgs.to(device)\n",
        "            targets = targets.to(device)\n",
        "            outputs = mymod(imgs)\n",
        "\n",
        "            loss = loss_fn(outputs, targets)\n",
        "            total_test_loss += loss.item()\n",
        "            # 统计在测试集上正确的个数，然后累加起来\n",
        "            accuracy = (outputs.argmax(1) == targets).sum()\n",
        "            total_accuracy += accuracy\n",
        "\n",
        "            if total_test_step % 1000 == 0:\n",
        "                writer.add_scalar(\"test_loss\", loss.item(), total_test_step)\n",
        "            total_test_step += 1\n",
        "\n",
        "    writer.add_scalar(\"test_accuracy\", total_accuracy / test_data_size, test_acc_step)\n",
        "    test_acc_step += 1\n",
        "\n",
        "torch.save(mymod, \"./mymod.pth\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
